#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <zlib.h>
#include "protocol.h"

// 计算CRC32校验和
uint32_t calculate_crc32(const void *data, size_t len)
{
    return crc32(0L, (const Bytef *)data, len);
}

// 发送完整的消息
int send_message(int sock, uint8_t type, uint16_t seq, const void *data, uint32_t len)
{
    if (sock < 0 || (data == NULL && len > 0))
    {
        return PROT_ERR_INVALID_PARAMS;
    }

    // 准备消息头
    message_header_t header;
    header.version = PROTOCOL_VERSION;
    header.type = type;
    header.seq = htons(seq);    // 网络字节序
    header.length = htonl(len); // 网络字节序

    // 计算CRC32校验和
    if (data && len > 0)
    {
        header.crc32 = htonl(calculate_crc32(data, len));
    }
    else
    {
        header.crc32 = 0;
    }

    // 发送消息头
    int sent = 0;
    int ret = 0;
    while (sent < sizeof(header))
    {
        ret = send(sock, ((char *)&header) + sent, sizeof(header) - sent, 0);
        if (ret <= 0)
        {
            if (ret < 0 && (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK))
            {
                continue; // 重试
            }
            perror("[SEND_MSG] Header send failed");
            return PROT_ERR_SEND_FAILED;
        }
        sent += ret;
    }

    // 发送消息体
    if (data && len > 0)
    {
        sent = 0;
        while (sent < len)
        {
            ret = send(sock, ((char *)data) + sent, len - sent, 0);
            if (ret <= 0)
            {
                if (ret < 0 && (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK))
                {
                    continue; // 重试
                }
                perror("[SEND_MSG] Body send failed");
                return PROT_ERR_SEND_FAILED;
            }
            sent += ret;
        }
    }
    return PROT_SUCCESS;
}

// 接收完整的消息
int recv_message(int sock, uint8_t *type, uint16_t *seq, void **data, uint32_t *len)
{
    if (sock < 0 || type == NULL || seq == NULL || data == NULL || len == NULL)
    {
        return PROT_ERR_INVALID_PARAMS;
    }

    *data = NULL; // 确保data指针初始为NULL
    *len = 0;

    message_header_t header;
    int received_total = 0;
    int ret = 0;
    int retry_count = 0;
    const int max_retries = 5; // 针对EAGAIN/EWOULDBLOCK的重试

    // 接收消息头
    while (received_total < sizeof(header))
    {
        ret = recv(sock, ((char *)&header) + received_total, sizeof(header) - received_total, 0);
        if (ret < 0)
        {
            if (errno == EINTR)
                continue;
            if (errno == EAGAIN || errno == EWOULDBLOCK)
            {
                retry_count++;
                if (retry_count >= max_retries)
                {
                    fprintf(stderr, "[RECV_MSG] Timeout waiting for header (EAGAIN/EWOULDBLOCK)\n");
                    return PROT_ERR_RECV_TIMEOUT_WOULD_BLOCK;
                }
                usleep(100000); // 100ms wait
                continue;
            }
            perror("[RECV_MSG] Header recv failed");
            return PROT_ERR_GENERIC_ERROR;
        }
        if (ret == 0)
        {
            fprintf(stderr, "[RECV_MSG] Connection closed by peer while receiving header.\n");
            return PROT_ERR_CONNECTION_CLOSED_BY_PEER;
        }
        received_total += ret;
    }

    // 检查协议版本
    if (header.version != PROTOCOL_VERSION)
    {
        fprintf(stderr, "[RECV_MSG] Protocol version mismatch: Expected %02X, Got %02X\n", PROTOCOL_VERSION, header.version);
        // 版本不匹配时，我们不知道消息体长度，无法安全地丢弃消息，连接可能已损坏
        return PROT_ERR_VERSION_MISMATCH;
    }

    // 转换网络字节序 (除了CRC，它在验证后转换)
    uint16_t received_seq = ntohs(header.seq);
    uint32_t expected_body_len = ntohl(header.length);
    uint32_t received_crc32 = ntohl(header.crc32); // CRC先转过来，因为它是对原始网络字节序的body计算的

    // 分配内存接收消息体
    if (expected_body_len > 0)
    {
        if (expected_body_len > MAX_MSG_BODY_LEN)
        { // 增加一个最大长度保护
            fprintf(stderr, "[RECV_MSG] Declared message body length %u exceeds max %u\n", expected_body_len, MAX_MSG_BODY_LEN);
            return PROT_ERR_MSG_TOO_SHORT; // 或者一个更特定的错误码如PROT_ERR_MSG_LEN_EXCEEDS_LIMIT
        }
        *data = malloc(expected_body_len);
        if (*data == NULL)
        {
            fprintf(stderr, "[RECV_MSG] Malloc failed for body (size: %u)\n", expected_body_len);
            return PROT_ERR_MALLOC_FAILED;
        }

        received_total = 0;
        retry_count = 0;
        while (received_total < expected_body_len)
        {
            ret = recv(sock, ((char *)(*data)) + received_total, expected_body_len - received_total, 0);
            if (ret < 0)
            {
                if (errno == EINTR)
                    continue;
                if (errno == EAGAIN || errno == EWOULDBLOCK)
                {
                    retry_count++;
                    if (retry_count >= max_retries)
                    {
                        fprintf(stderr, "[RECV_MSG] Timeout waiting for body (EAGAIN/EWOULDBLOCK)\n");
                        free(*data);
                        *data = NULL;
                        return PROT_ERR_RECV_TIMEOUT_WOULD_BLOCK;
                    }
                    usleep(100000); // 100ms wait
                    continue;
                }
                perror("[RECV_MSG] Body recv failed");
                free(*data);
                *data = NULL;
                return PROT_ERR_GENERIC_ERROR;
            }
            if (ret == 0)
            {
                fprintf(stderr, "[RECV_MSG] Connection closed by peer while receiving body.\n");
                free(*data);
                *data = NULL;
                return PROT_ERR_CONNECTION_CLOSED_BY_PEER;
            }
            received_total += ret;
        }

        // 验证CRC32校验和
        uint32_t calc_crc32 = calculate_crc32(*data, expected_body_len);
        if (calc_crc32 != received_crc32)
        {
            fprintf(stderr, "[RECV_MSG] CRC32 mismatch: Calculated %08X, Expected %08X\n", calc_crc32, received_crc32);
            free(*data);
            *data = NULL;
            return PROT_ERR_CRC_MISMATCH;
        }
    }
    else
    {
        // 长度为0的消息，CRC也应该为0
        if (received_crc32 != 0)
        {
            fprintf(stderr, "[RECV_MSG] CRC32 mismatch for zero-length body: Expected 0, Got %08X\n", received_crc32);
            return PROT_ERR_CRC_MISMATCH;
        }
        *data = NULL; // 确保data为NULL如果长度为0
    }

    *type = header.type;
    *seq = received_seq;
    *len = expected_body_len;

    return PROT_SUCCESS;
}

// 接收特定类型的消息 (可选匹配序列号)
int recv_specific_message(int sock, uint8_t expected_type, uint16_t expected_seq, void **data, uint32_t *len)
{
    if (data == NULL || len == NULL)
        return PROT_ERR_INVALID_PARAMS;

    uint8_t type_recv;
    uint16_t seq_recv;
    void *data_recv = NULL;
    uint32_t len_recv = 0;

    int ret = recv_message(sock, &type_recv, &seq_recv, &data_recv, &len_recv);
    if (ret != PROT_SUCCESS)
    {
        // recv_message 已经打印了错误, data_recv 应该为 NULL
        if (data_recv)
        {
            free(data_recv);
            data_recv = NULL;
        } // 以防万一
        *data = NULL;
        *len = 0;
        return ret; // 直接返回recv_message的错误码
    }

    // 首先检查是否是服务器发送的通用错误消息
    if (type_recv == MSG_TYPE_ERROR)
    {
        fprintf(stderr, "[SPECIFIC_MSG] Received MSG_TYPE_ERROR from server.\n");
        if (data_recv && len_recv >= sizeof(error_message_t))
        {
            error_message_t *err_msg = (error_message_t *)data_recv;
            fprintf(stderr, "[SPECIFIC_MSG] Server Error Code: %u, Message: %s\n", err_msg->code, err_msg->message);
            // 可选：将 err_msg->code 存放到一个会话相关的变量中供上层获取
        }
        else if (data_recv)
        {
            fprintf(stderr, "[SPECIFIC_MSG] Received MSG_TYPE_ERROR with invalid body length: %u\n", len_recv);
        }
        if (data_recv)
        {
            free(data_recv);
            data_recv = NULL;
        }
        *data = NULL;
        *len = 0;
        return PROT_ERR_SERVER_SENT_ERROR_MSG; // 返回特定错误码表示服务器报告了错误
    }

    // 检查消息类型是否匹配
    if (type_recv != expected_type)
    {
        fprintf(stderr, "[SPECIFIC_MSG] Unexpected message type: Expected %02X, Got %02X. Seq_exp=%u, Seq_got=%u\n",
                expected_type, type_recv, expected_seq, seq_recv);
        if (data_recv)
        {
            free(data_recv);
            data_recv = NULL;
        }
        *data = NULL;
        *len = 0;
        // 应该丢弃这个不期望的消息体，但这里简单返回错误
        return PROT_ERR_UNEXPECTED_MSG_TYPE;
    }

    // 如果需要，检查序列号 (expected_seq == 0 表示不检查序列号)
    if (expected_seq != 0 && seq_recv != expected_seq)
    {
        fprintf(stderr, "[SPECIFIC_MSG] Sequence number mismatch: Expected %u, Got %u for Type %02X\n",
                expected_seq, seq_recv, expected_type);
        if (data_recv)
        {
            free(data_recv);
            data_recv = NULL;
        }
        *data = NULL;
        *len = 0;
        return PROT_ERR_SEQ_NUM_MISMATCH;
    }

    // 所有检查通过
    *data = data_recv;
    *len = len_recv;
    return PROT_SUCCESS;
}

// 发送错误消息
int send_error_message(int sock, uint16_t seq, uint8_t error_code, const char *error_msg)
{
    error_message_t error;
    error.code = error_code;

    if (error_msg)
    {
        strncpy(error.message, error_msg, sizeof(error.message) - 1);
        error.message[sizeof(error.message) - 1] = '\0';
    }
    else
    {
        error.message[0] = '\0';
    }

    return send_message(sock, MSG_TYPE_ERROR, seq, &error, sizeof(error));
}